import numpy as np
import matplotlib.pyplot as pt
import numpy.linalg as la
import scipy.sparse.linalg as lin
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import griddata
import seaborn as sns
%matplotlib inline
# %matplotlib qt

# When numerically solving the structure of hydrogen atom using finite difference method, 
# the main objective is to establish 1-1 correspondence between the col vector which represents coordinates:
# (000) (001)....(NNN), with the horizontal location of matrix (0,1...N^3) to put values in 
# matrix where the numbers should be calculated. 
# if we focus on certain point (x,y,z), it has a location in the col vector, and its neigbourhood(x-1,y,z)
# has various positions in the horizontal row in the matrix, and we need to find the positions, that's why we
# build the identification-table. 

L = 31 # dimension of the box
N = 20 # Accuracy; even number
location = []
for i in range(0,N):
    for j in range(0,N):
        for k in range(0,N):
            location.append((i,j,k)) # (000),001,...00N, 011,...0N1,...0NN, 100, 101....NNN
            
def matrix_to_index(a,b,c):
    return location.index((a,b,c)) # e.g (0,0,3)-> 3
def index_to_ijk(i):
    return location[i] #e.g 5->(0,0,5)

X = np.linspace(-L/2,L/2,N)
Y = np.linspace(-L/2,L/2,N)
Z = np.linspace(-L/2,L/2,N)
d = X[1]-X[0]

# Constructing Hamiltonian: -p^2/2m + V(1/r)

H=np.zeros((N**3,N**3))

for i in range(0,N**3): # loop all xyz combination
    x = index_to_ijk(i)[0]
    y = index_to_ijk(i)[1]
    z = index_to_ijk(i)[2]
    
    H[i,matrix_to_index(x,y,z)] = -6   # f''(x) = f''(x+d)-2f(x)+f''(x-d) / d^2 and sum! 
    if x+1<N:
        H[i,matrix_to_index(x+1,y,z)] = 1
    if x-1>=0:
        H[i,matrix_to_index(x-1,y,z)] = 1
    if y+1<N:
        H[i,matrix_to_index(x,y+1,z)] = 1
    if y-1>=0:
        H[i,matrix_to_index(x,y-1,z)] = 1
    if z+1<N:
        H[i,matrix_to_index(x,y,z+1)] = 1
    if z-1>=0:
        H[i,matrix_to_index(x,y,z-1)] = 1
H = -H/d**2

for i in range(0,N**3):
    x = index_to_ijk(i)[0]
    y = index_to_ijk(i)[1]
    z = index_to_ijk(i)[2]
    H[i,i] += -2/np.sqrt(X[x]**2+Y[y]**2+Z[z]**2)

sns.set(font_scale=1.5)
sns.set_style("whitegrid")
pt.figure(figsize=(5,5))
pt.imshow(H[0:50,0:50],cmap='gray_r',interpolation='none')
pt.colorbar()
pt.show()

# Solving for eigensystem using spars matrix technique

vals, vecs = lin.eigsh(H,k=20,which='SA') # first k states with lowest energies

## Different energy states 
sns.set(font_scale=1.5)
sns.set_style("whitegrid")

for n in range(0,20):    # including ground states
    fig = pt.figure(figsize=(15,6))
    
    # Density distribution across xy plane

    ax = fig.add_subplot(121, projection='3d')
    ax.tick_params(axis='x', pad=1)
    ax.tick_params(axis='y', pad=1)
    ax.tick_params(axis='z', pad=1)
    xs=[]
    ys=[]
    rho=[]
    for i in range(0,len(vecs[:,0])):
        if index_to_ijk(i)[2]==0:
            xs.append(X[index_to_ijk(i)[0]])
            ys.append(Y[index_to_ijk(i)[1]])
            rho.append(vecs[:,n][i]**2*10**5) # *1000 to make things look nicer
            
    # interpolate
    xi = yi = np.arange(-L/2,L/2,1)
    xi,yi = np.meshgrid(xi,yi)
    zi = griddata((xs,ys),rho,(xi,yi),method='cubic')
    pt.contourf(xi,yi,zi,np.arange(0,np.max(rho),np.max(rho)/200),cmap='rainbow',alpha=0.3)
        
    ax.scatter(xs, ys, rho, c=rho, marker='o',cmap='rainbow')
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    pt.title('Density Distrbution across a flat surface')
    
    # Spacial density distribution

    ax = fig.add_subplot(122, projection='3d')
    ax.tick_params(axis='x', pad=1)
    ax.tick_params(axis='y', pad=1)
    ax.tick_params(axis='z', pad=1)
    xs=[]
    ys=[]
    zs=[]
    rho=[]
    for i in range(0,len(vecs[:,0])):
        xs = np.append(xs, X[index_to_ijk(i)[0]])
        ys = np.append(ys, Y[index_to_ijk(i)[1]])
        zs = np.append(zs, Z[index_to_ijk(i)[2]])
        rho = np.append(rho, vecs[:,n][i]**2)
    rho=np.ma.masked_where(rho < 0.0003, rho)
    
    ax.scatter(xs, ys, zs, c=rho, marker='.',alpha=1,cmap='rainbow')
    ax.set_xlim(-14,14)
    ax.set_ylim(-14,14)
    ax.set_zlim(-14,14)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    pt.title(str(n)+'th Excited State Density Distribution ($\Psi^2$)')
    pt.show()